package com.pg.agent.engine.service.utils;


import cn.hutool.core.bean.BeanUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.pg.agent.engine.client.component.APIComponent;
import com.pg.agent.engine.client.dto.SOPCallVO;
import com.pg.agent.engine.client.message.InPutMessage;
import com.pg.agent.engine.client.message.OutPutMessage;
import com.pg.agent.engine.client.model.ComponentOutParam;
import com.pg.agent.engine.client.model.agent.ComponentParam;
import com.pg.agent.engine.client.model.llm.CallLog;
import com.pg.agent.engine.client.vo.HttpParam;
import com.pg.agent.engine.client.vo.SOPResultVO;
import com.pg.base.enums.ResultEnum;
import com.pg.base.exception.BizException;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
@Data
public class ComponentSSEListener extends EventSourceListener {

    private CountDownLatch countDownLatch = new CountDownLatch(1);

    private HttpParam httpParam;

    private String messageId;

    private HttpServletResponse rp;

    private OutPutMessage outMessage;

    private InPutMessage message;

    private boolean isFirst = true;

    private Integer code = 0;

    private String msg = "success";

    private JSONObject jsonValue;

    private StringBuffer output;

    private APIComponent component;

    private boolean isReturn = true;

    public ComponentSSEListener(HttpParam httpParam, InPutMessage message, HttpServletResponse response, String messageId, OutPutMessage outMessage, APIComponent component) {
        this.httpParam = httpParam;
        this.rp = response;
        this.message = message;
        this.messageId = messageId;
        this.outMessage = outMessage;
        this.output = new StringBuffer();
        this.component = component;
    }

    /**
     * {@inheritDoc}
     * 建立sse连接
     */
    @Override
    public void onOpen(final EventSource eventSource, final Response response) {
        if (rp != null) {
            rp.setContentType("text/event-stream");
            rp.setCharacterEncoding("UTF-8");
            rp.setStatus(200);
            log.info("建立sse连接..." + JSON.toJSONString(httpParam));
            outMessage.setTimeLong(System.currentTimeMillis() - outMessage.getCallStartTime().getTime());
            if (message != null && message.isDebug()) {
                outMessage.setCallLogList(BeanUtil.copyToList(message.getCallLogList(), CallLog.class));
            }
            SOPResultVO vo = SOPResultVO.instance("", false, messageId, outMessage);
            if (outMessage.isFirst()) {
                outMessage.setFirst(false);
                vo.setFirst(1);
            }
            try {
                rp.getWriter().write("data:" + JSONObject.toJSONString(vo) + "\n\n");
                rp.getWriter().flush();
            } catch (IOException e) {
                setSuccessStatus(-1, e.getMessage());
                log.error("消息错误[" + JSON.toJSONString(httpParam) + "]", e);
                countDownLatch.countDown();
            }

        } else {
            log.info("客户端非sse推送" + JSON.toJSONString(httpParam));
        }
    }

    /**
     * 事件
     *
     * @param eventSource
     * @param id
     * @param type
     * @param data
     */
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        try {
            Map<String, AtomicBoolean> atomicIntegerMap = SessionConnectManager.getInstance().getStoryMap();
            AtomicBoolean atomicBoolean = atomicIntegerMap.get(message.getSessionId());
            if (atomicBoolean != null && !atomicBoolean.get()) {
                SessionConnectManager.getInstance().addStopSessionId(message.getSessionId());
                log.info("流程中断");
                countDownLatch.countDown();
            }
            boolean finish = false;
            String content = null;
            try {
                if (isFirst) {
                    if (JSON.isValidObject(data)) {
                        jsonValue = JSON.parseObject(data);
                    } else {
                        jsonValue = new JSONObject();
                        jsonValue.put("data", data);
                    }

                    isFirst = false;
                }
                log.info("data:{}", data);
                JSONObject result = JSONObject.parseObject(data);
                content = result.getJSONObject("data").getString("content");
                if (StringUtils.isNotEmpty(content)) {
                    output.append(content);
                }
            } catch (Exception e) {
                content = data;
            }

            if (rp != null) {
                if ("\n".equals(data)) {
                    rp.getWriter().write("event:" + type + "\n");
                    rp.getWriter().write("id:" + id + "\n");
                    rp.getWriter().write("data:\n\n");
                    rp.getWriter().flush();
                } else {
                    SOPResultVO vo = null;
                    if (jsonValue != null && isReturn) {
                        log.info("=============第一次进入===================");
                        isReturn = false;

                        SOPCallVO sopCallVO1 = new SOPCallVO();
                        String result = message.getAgent().getResponse();
                        log.info("result:{}", result);
                        for (ComponentParam componentParam : message.getAgent().getComponentParamList()) {
                            Long componentId = componentParam.getComponentId();
                            String paramCode = componentParam.getParamCode();
                            List<ComponentOutParam> list = JSONObject.parseArray(component.getOutParam(), ComponentOutParam.class);
                            List<ComponentOutParam> allList = new ArrayList<>();
                            JSONUtils.getAllComponentOutParam(list, allList, "");
                            JSONObject jsonObject = new JSONObject();
                            Map<String, Object> flattenedMap = new LinkedHashMap<>();
                            JSONUtils.flattenJson("root", jsonObject, flattenedMap);
                            flattenedMap.put("root", jsonValue);
                            for (ComponentOutParam outParam : allList) {
                                //前缀没有root,如果需要获取值，需要去除前缀root
                                jsonObject.put(outParam.getNewParamCode(), flattenedMap.get(outParam.getNewParamCode()));
                            }
                            try {
                                if (jsonObject == null || componentId.longValue() != component.getId().longValue()) {
                                    result = result.replace("{comp-" + componentId + "-" + paramCode + "}", "null");
                                } else {
                                    result = result.replace("{comp-" + componentId + "-" + paramCode + "}",
                                            jsonObject.get(paramCode) == null ? "null" : JSON.toJSONString(jsonObject.get(paramCode), SerializerFeature.DisableCircularReferenceDetect));
                                }
                            } catch (Exception e) {
                                log.info("解析异常：{}", e.getMessage());
                            }
                        }
                        sopCallVO1.setAgentResponse(result);
                        vo = SOPResultVO.instance(null, finish, messageId, sopCallVO1);
                        log.info("vo:{}", JSONObject.toJSONString(vo));
                    } else {
                        vo = SOPResultVO.instance(content, finish, messageId, null);
                        log.info("vo1:{}", JSONObject.toJSONString(vo));
                    }
                    rp.getWriter().write("data:" + JSONObject.toJSONString(vo) + "\n\n");
                    rp.getWriter().flush();
                }
            }
        } catch (Exception e) {
            log.info("data:{}", data);
            log.error("消息错误[" + JSON.toJSONString(httpParam) + "]", e);
            setSuccessStatus(-1, e.getMessage());
            countDownLatch.countDown();

            throw new BizException(ResultEnum.ERROR, e);
        }

    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onClosed(final EventSource eventSource) {
        //     log.info("sse连接关闭:{}", param.getMessageId());
        log.info("结果输出:{}", output.toString());
        countDownLatch.countDown();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onFailure(final EventSource eventSource, final Throwable t, final Response response) {
        ResponseBody body = response.body();
        try {
            String value = body.string();
            if (value != null) {
                jsonValue = JSONObject.parseObject(value);
            }
            setSuccessStatus(1000, "执行组件报错");
        } catch (IOException e) {
            log.info("读取失败");
            setSuccessStatus(-1, "执行组件出错");
        }
        //log.error("使用事件源时出现异常... [响应：{}]...", t.getMessage());

        countDownLatch.countDown();
    }

    public JSONObject getJsonValue() {
        if (jsonValue == null) {
            jsonValue = new JSONObject();
        }
        if (jsonValue.getInteger("code") != null && jsonValue.getInteger("code") == 0) {
            jsonValue.getJSONObject("data").put("content", output.toString());
        }
        return jsonValue;
    }

    public CountDownLatch getCountDownLatch() {
        return this.countDownLatch;
    }

    private void setSuccessStatus(Integer code, String msg) {
        this.code = code;
        this.msg = msg;
    }


}

